"""
对结果进行排序
"""
import torch
import json

import config
from dnn.sort.siamese import SiameseNetWork


class DnnSort:
    def __init__(self):
        # 加载模型
        self.model = SiameseNetWork().to(config.device)
        self.model.load_state_dict(torch.load(config.sort_save_model_path))

        # 加载存储的问题文件数据
        self.qa_dict = json.load(open(config.recall_corpus_tfidf_path, mode='r', encoding="UTF-8"))
        pass

    def sort(self, sentence, recall_list):
        """
        sentence: 用户输入的句子： {"cut": str, "cut_by_word": str, "entity": list}
        recall_list: 对用户输入的句子进行召回得到的结果
        """
        # 1. 构建模型的输入
        input1 = [[sentence["cut"]]] * len(recall_list)
        # print(input1)
        input2 = [self.qa_dict[i]["cut_by_word"] for i in recall_list]

        # 2. 对输入数据进行处理
        input1 = torch.LongTensor([config.sort_q_ws_model.transform(i,
                                                                    max_length=config.sort_max_len)
                                   for i in input1]).to(config.device)
        input2 = torch.LongTensor([config.sort_sim_q_ws_model.transform(i,
                                                                        max_length=config.sort_max_len)
                                   for i in input2]).to(config.device)

        # 3. 模型处理返回结果, 计算用户输入的句子和召回句子的相似度
        # output: [len(recall_list), 2]
        output = self.model(input1, input2)
        # ouptput: [len(recall_list)]
        output = torch.nn.functional.softmax(output, dim=-1)
        output = output[:, -1].squeeze(-1)
        # output = output.max(dim=-1)[-1]

        # 4. 对3处理的结果进行排序
        print(list(zip(recall_list, output.detach().cpu().numpy())))
        best_q, best_prob = sorted(list(zip(recall_list, output.cpu().detach().numpy())),
                                   key=lambda x: x[-1], reverse=True)[0]

        # 5. 返回最佳结果
        if best_prob > 0.7:
            return best_q
        else:
            return "无法回答"
        pass
